#!/bin/bash
# Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

############################################################################
## This script generates the model repository needed for TensorRT testing
## testing on the Jetson Xavier. Generating these models requires having
## TensorRT on the system.
##
## 1. Update Jetpack version to match what is being used by the
## tritonserver release being tested.
##
## 2. Set CUDA_DEVICE to the ID of the CUDA device present on the
## system that you want to target for the generated models.
##
## 3. Run this script to create a '/tmp/<TRITON_VERSION>_xavier' directory
## containing all the TensorRT models needed for CI testing.
##
## 4. Copy all non-TensorRT models from the dlcluster directory
## '/mnt/dldata/inferenceserver/<TRITON_VERSION>' to
## '/mnt/dldata/inferenceserver/<TRITON_VERSION>_xavier'.
##
## 5. Copy the generated models from '/tmp/<TRITON_VERSION>_xavier' to the
## dlcluster directory '/mnt/dldata/inferenceserver/<TRITON_VERSION>_xavier'.
##
## 6. Remember to delete the '/tmp/<TRITON_VERSION>_xavier' directory after.
##
############################################################################

TRITON_VERSION=${TRITON_VERSION:=23.04}
CUDA_DEVICE=${RUNNER_ID:=0}

HOST_BUILD_DIR=${HOST_BUILD_DIR:=/tmp}
TMP_DIR=$HOST_BUILD_DIR/${TRITON_VERSION}

SRCDIR=$HOST_BUILD_DIR/gen_srcdir
DESTDIR=$TMP_DIR/qa_model_repository

DATADEPENDENTDIR=$TMP_DIR/qa_trt_data_dependent_model_repository
DYNASEQDESTDIR=$TMP_DIR/qa_dyna_sequence_model_repository
DYNASEQIMPLICITDESTDIR=$TMP_DIR/qa_dyna_sequence_implicit_model_repository
FORMATDESTDIR=$TMP_DIR/qa_trt_format_model_repository
IDENTITYBIGDESTDIR=$TMP_DIR/qa_identity_big_model_repository
IDENTITYDESTDIR=$TMP_DIR/qa_identity_model_repository
IMPLICITSEQDESTDIR=$TMP_DIR/qa_sequence_implicit_model_repository
RAGGEDDESTDIR=$TMP_DIR/qa_ragged_model_repository
RESHAPEDESTDIR=$TMP_DIR/qa_reshape_model_repository
SEQDESTDIR=$TMP_DIR/qa_sequence_model_repository
SHAPEDESTDIR=$TMP_DIR/qa_shapetensor_model_repository
VARDESTDIR=$TMP_DIR/qa_variable_model_repository
VARIMPLICITSEQDESTDIR=$TMP_DIR/qa_variable_sequence_implicit_model_repository
VARSEQDESTDIR=$TMP_DIR/qa_variable_sequence_model_repository


rm -fr $DESTDIR
rm -fr $DATADEPENDENTDIR
rm -fr $DYNASEQDESTDIR
rm -fr $DYNASEQIMPLICITDESTDIR
rm -fr $FORMATDESTDIR
rm -fr $IDENTITYBIGDESTDIR
rm -fr $IDENTITYDESTDIR
rm -fr $IMPLICITSEQDESTDIR
rm -fr $RAGGEDDESTDIR
rm -fr $RESHAPEDESTDIR
rm -fr $SEQDESTDIR
rm -fr $SHAPEDESTDIR
rm -fr $VARDESTDIR
rm -fr $VARIMPLICITSEQDESTDIR
rm -fr $VARSEQDESTDIR

mkdir -p $SRCDIR
mkdir -p $DESTDIR

mkdir -p $DATADEPENDENTDIR
mkdir -p $DYNASEQDESTDIR
mkdir -p $DYNASEQIMPLICITDESTDIR
mkdir -p $FORMATDESTDIR
mkdir -p $IDENTITYBIGDESTDIR
mkdir -p $IDENTITYDESTDIR
mkdir -p $IMPLICITSEQDESTDIR
mkdir -p $RAGGEDDESTDIR
mkdir -p $RESHAPEDESTDIR
mkdir -p $SEQDESTDIR
mkdir -p $SHAPEDESTDIR
mkdir -p $VARDESTDIR
mkdir -p $VARIMPLICITSEQDESTDIR
mkdir -p $VARSEQDESTDIR

cp ./gen_ensemble_model_utils.py $SRCDIR/.
cp ./gen_qa_dyna_sequence_implicit_models.py $SRCDIR/.
cp ./gen_qa_dyna_sequence_models.py $SRCDIR/.
cp ./gen_qa_identity_models.py $SRCDIR/.
cp ./gen_qa_implicit_models.py $SRCDIR/.
cp ./gen_qa_models.py $SRCDIR/.
cp ./gen_qa_noshape_models.py $SRCDIR/.
cp ./gen_qa_ragged_models.py $SRCDIR/.
cp ./gen_qa_reshape_models.py $SRCDIR/.
cp ./gen_qa_sequence_models.py $SRCDIR/.
cp ./gen_qa_trt_data_dependent_shape.py $SRCDIR/.
cp ./gen_qa_trt_format_models.py $SRCDIR/.
cp ./gen_qa_trt_plugin_models.py $SRCDIR/.
cp ./test_util.py $SRCDIR/.

# TensorRT
set -e
export TRT_SUPPRESS_DEPRECATION_WARNINGS=1
cd $SRCDIR
# Models using shape tensor i/o
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt-shape-io --models_dir=$SHAPEDESTDIR
chmod -R 777 $SHAPEDESTDIR
python3 $SRCDIR/gen_qa_models.py --tensorrt --models_dir=$DESTDIR
chmod -R 777 $DESTDIR
python3 $SRCDIR/gen_qa_models.py --tensorrt --variable --models_dir=$VARDESTDIR
chmod -R 777 $VARDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt --models_dir=$IDENTITYDESTDIR
chmod -R 777 $IDENTITYDESTDIR
python3 $SRCDIR/gen_qa_identity_models.py --tensorrt-big --models_dir=$IDENTITYBIGDESTDIR
chmod -R 777 $IDENTITYBIGDESTDIR
python3 $SRCDIR/gen_qa_reshape_models.py --tensorrt --variable --models_dir=$RESHAPEDESTDIR
chmod -R 777 $RESHAPEDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt --models_dir=$SEQDESTDIR
chmod -R 777 $SEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --tensorrt --models_dir=$IMPLICITSEQDESTDIR
chmod -R 777 $IMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_implicit_models.py --tensorrt --variable --models_dir=$VARIMPLICITSEQDESTDIR
chmod -R 777 $VARIMPLICITSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt --models_dir=$DYNASEQDESTDIR
chmod -R 777 $DYNASEQDESTDIR
python3 $SRCDIR/gen_qa_sequence_models.py --tensorrt --variable --models_dir=$VARSEQDESTDIR
chmod -R 777 $VARSEQDESTDIR
python3 $SRCDIR/gen_qa_dyna_sequence_implicit_models.py --tensorrt --models_dir=$DYNASEQIMPLICITDESTDIR
chmod -R 777 $DYNASEQIMPLICITDESTDIR
python3 $SRCDIR/gen_qa_ragged_models.py --tensorrt --models_dir=$RAGGEDDESTDIR
chmod -R 777 $RAGGEDDESTDIR
python3 $SRCDIR/gen_qa_trt_format_models.py --models_dir=$FORMATDESTDIR
chmod -R 777 $FORMATDESTDIR
python3 $SRCDIR/gen_qa_trt_data_dependent_shape.py --models_dir=$DATADEPENDENTDIR
chmod -R 777 $DATADEPENDENTDIR
